"""
手写体数字识别模型的训练和保存
数据来源：MNIST数据库（tensorflow.examples.tutorials.mnist）
        Modified National Institute of Standards and Technology database
        修改后的美国国家标准与技术研究所数据库
训练集：mnist.train
测试集：mnist.test
每个样本：28×28=784的灰度图片
cnn结构：LeNet Yann LeCun在1998年提出，用于解决手写数字识别的视觉任务。
        1. input输入层（将原始输入转为n×32*32*1的形状）
        2. c1卷积层（5*5卷积核，单通道，6输出 → n×28*28*6）
        3. s2池化层（2*2average池化 → relu → n×14*14*6）
        4. c3卷积层（6+6+3+1=16输出 → n×10*10*16）
            6个（5*5卷积核，3通道，1输出）
            6个（5*5卷积核，4通道，1输出）
            3个（5*5卷积核，(2+2)通道，1输出）
            1个（5*5卷积核，6通道，1输出）
        5. s4池化层（2*2average池化 → relu → n×5*5*16）
        6. c5卷积层（5*5卷积核，单通道，120输出 → n×1*1*120 → n*120）
        7. f6全连接层（84个神经元 → n×84）
        8. output全连接层（10个神经元 → 极差标准化 → n×10）
"""
import os
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data


class LeNetCNN:
    def __init__(self, data_dir='', model_dir=None):
        """
        构造方法
        :param data_dir: 数据集所在目录，默认从网络下载
        :param model_dir: 模型的存储目录，为None时则不保存模型
        """
        if type(model_dir) == str:
            if not model_dir.endswith('/'):
                raise Exception('model_dir必须以“/”结尾')

        self.current_dir = os.path.dirname(__file__)
        self.model_dir = model_dir and os.path.join(self.current_dir, model_dir)
        self.data_dir = os.path.join(self.current_dir, data_dir)
        self.sess = tf.Session()
        self.data = input_data.read_data_sets(self.data_dir, one_hot=True)
        # print(self.data.train.next_batch(1))

        # 原始输入(真实数据)
        self.input = tf.placeholder(tf.float32, [None, 784], name="input_placeholder")
        # print(self.input)  # shape=(?, 784)

        # 原始输入对应的标签(真实结果)
        self.label = tf.placeholder(tf.float32, [None, 10], name="label_placeholder")
        # print(self.label)

        # 神经网络模型(结果)
        self.result = self.create_lenet_cnn()
        # print(self.result)

        # 模型准确率
        self.accuracy = self.calc_accuracy()

        # 损失值
        self.loss = self.create_loss()

        # 训练步骤
        self.train_step = self.create_train_step()

        # 模型状态保存器
        var_list = [v for v in tf.global_variables() if 'Adam' not in v.name and 'Variable' not in v.name]
        self.saver = tf.train.Saver(var_list)

        # 变量初始化操作
        self.init_op = tf.global_variables_initializer()

        # 模型变量初始化
        self.init_vars()

    def create_lenet_cnn(self):
        """
        创建LeNet卷积神经网络
        :return: 返回网络输出结果
        """
        # 输入层
        input = self.create_input(self.input)

        # c1卷积层
        c1 = self.create_c1(input)

        # s2池化层-下采样
        s2 = self.create_s2(c1)

        # c3卷积层
        c3 = self.create_c3(s2)

        # s4池化层-下采样
        s4 = self.create_s4(c3)

        # c5卷积层
        c5 = self.create_c5(s4)

        # f6全连接层
        f6 = self.create_f6(c5)

        # output全连接层
        output = self.create_output(f6)

        return output

    def create_fc(self, input, input_features, out_features):
        """
        创建全连接层
        :param input: 输入（经过拉伸的一维张量）
        :param input_features: 输入特征值数量
        :param out_features: 输出特征值数量（神经元数量）
        :return:
        """
        # 初始化权重
        weight = self.init_weight([input_features, out_features], 'fc_weight')

        # 初始化偏置，每个输出(神经元)对应一个偏置和n个权重
        bias = self.init_bias([out_features], 'fc_bias')

        # 神经网络计算
        fc = tf.matmul(input, weight) + bias

        return fc

    def init_weight(self, shape, name):
        """
        初始化权重(卷积核可以理解为特殊的权重)
        :param shape: 权重张量形状
        :return: 初始化后的权重值
        """
        initial = tf.random_normal(shape)
        return tf.Variable(initial, name=name)

    def init_bias(self, shape, name):
        """
        初始化偏置
        :param shape: 形状
        :return: 初始化后的偏置值
        """
        initial = tf.zeros(shape)
        return tf.Variable(initial, name=name)

    def calc_accuracy(self):
        """
        计算模型准确率
        :return: 返回模型准确率
        """
        # 评价
        correct_prediction = tf.equal(
            tf.argmax(self.result, 1),  # 取出预测概率中最大的值的索引
            tf.argmax(self.label, 1)  # 取出真实概率中最大的值的索引
        )

        # 将上一步得到的bool类型数组转换为浮点型，并求准确率
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        return accuracy

    def create_loss(self):
        """
        创建损失函数
        :return: 返回损失值
        """
        # 交叉熵
        cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
            labels=self.label,  # 真实值
            logits=self.result  # 预测值
        )

        # 损失值
        loss = tf.reduce_mean(cross_entropy)

        return loss

    def create_train_step(self):
        """
        创建训练操作
        :return: 返回训练操作
        """
        # 优化器
        optimizer = tf.train.AdamOptimizer(0.001)

        # 步骤
        train_step = optimizer.minimize(self.loss)

        return train_step

    def init_vars(self, restore=True):
        """
        模型全局变量初始化
        :param restore: 是否从保存在模型目中的checkpoint恢复全局变量
        :return:
        """
        # vars = tf.global_variables()
        # for var in vars:
        #     print(var)
        self.sess.run(self.init_op)
        if restore and self.model_dir:
            checkpoint = os.path.join(self.model_dir, 'checkpoint')
            if os.path.exists(checkpoint):
                self.saver.restore(self.sess, self.model_dir)
            else:
                print('模型文件路径有误')

    def save_vars(self):
        """
        保存模型变量
        :return:
        """
        if self.model_dir:
            self.saver.save(self.sess, self.model_dir)
            print('save model success')

    def get_result(self, fetches, input, label=None):
        """
        获取模型计算结果
        :param fetches: 待获取的结果
        :param input: 输入n*784
        :param label: 输入对应的标签n*10
        :return: 返回计算的结果
        """
        params = {self.input: input}
        if label is not None:
            params[self.label] = label

        return self.sess.run(fetches, params)

    def train(self):
        """
        执行训练操作
        :return:
        """
        epochs = 20  # 训练轮次
        batch_size = 100  # 批次大小

        for i in range(epochs):  # 迭代训练
            total_batch = int(self.data.train.num_examples / batch_size)  # 计算批次数量，int为下取整
            for j in range(total_batch):
                batch = self.data.train.next_batch(batch_size)  # 获取一个批次样本
                t, acc, loss = self.get_result(
                    [self.train_step, self.accuracy, self.loss],
                    batch[0],
                    batch[1]
                )

                if j % 100 == 0:
                    print("epoch: %d, pass: %d, acc: %f, loss: %f" % (i, j, acc, loss))

    def eval(self, input, label):
        """
        模型测试评估
        :param input: 输入（测试数据）
        :param label: 标签（真实结果）
        :return: 返回测试准确率
        """
        test_acc = self.get_result(self.accuracy, input, label)
        print('Test accuracy %f' % test_acc)
        return test_acc

    def close(self):
        self.sess.close()

    def recognize(self, image):
        """
        识别一张数字手写体图片（白底黑字）
        :param image: 能被cv2识别的图片数据
        :return: 识别的结果
        """
        img = image
        if len(image.shape) == 3:
            img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        img = cv2.resize(img, (28, 28))
        img = 1 - img / 255
        # cv2.imshow('img',img)
        # cv2.waitKey()
        # cv2.destroyAllWindows()
        params = {self.input: img.reshape((1, 784))}
        result = self.sess.run(self.result, params)
        result = np.argmax(result[0])
        return result

    @staticmethod
    def channels_slice(input, start, size):
        """
        多通道特征图截取
        :param input: 待截取的张量n*Height*Width*Channels
        :param start: 从第几张特征图开始
        :param size: 截取几张特征图（头尾相连）
        :return: 返回截取的张量n*Height*Width*size
        """
        channels = input.shape[-1]
        end = start + size
        if end > channels:
            end -= channels
            a = input[:, :, :, start:]
            b = input[:, :, :, :end]
            return tf.concat((a, b), 3)
        else:
            return input[:, :, :, start:end]

    def create_conv(self, filter_shape, input):
        """
        卷积运算
        :param filter_shape: 卷积核形状，HWIO格式[filter_height, filter_width, in_channels, out_channels]
        :param input: 输入，待卷积的图片，NHWC格式
        :return: 返回卷积结果
        """
        filter = self.init_weight(filter_shape, 'conv_weight')
        bias = self.init_bias(filter_shape[-1:], 'conv_bias')
        conv = tf.nn.conv2d(
            input,  # 输入 NHWC格式
            filter,  # HWIO格式[filter_height, filter_width, in_channels, out_channels]
            strides=[1, 1, 1, 1],  # 步长 [N, H, W, C]
            padding='VALID'
        )

        return conv + bias

    def create_pool(self, input):
        """
        池化运算
        :param input: 输入，待池化的图片，NHWC格式
        :return: 返回池化的结果
        """
        pool = tf.nn.avg_pool(
            input,  # 输入
            ksize=[1, 2, 2, 1],  # [N,H,W,C]
            strides=[1, 2, 2, 1],  # [N,H,W,C]
            padding='VALID'
        )

        return tf.nn.relu(pool)

    def create_input(self, input):
        """
        创建LeNet第1层：输入层（将原始输入转为n×32*32*1的形状）
        :param input: 原始输入n×784
        :return: 返回n×32*32*1的张量
        """
        t = tf.reshape(input, [-1, 28, 28, 1])  # NHWC格式，方便使用tf做卷积运算
        output = tf.image.resize(t, [32, 32])
        # paddings = tf.constant([[0, 0], [2, 2], [2, 2], [0, 0]])
        # output = tf.pad(t, paddings)
        # print(output)  # shape=(?, 32, 32, 1)
        return output

    def create_c1(self, input):
        """
        创建LeNet第2层：卷积层（5*5卷积核，单通道，6输出）
        :param input: 第1层的输出n×32*32*1
        :return: 返回n×28*28*6的张量
        """
        output = self.create_conv([5, 5, 1, 6], input)
        # print(output)  # shape=(?, 28, 28, 6)
        return output

    def create_s2(self, input):
        """
        创建LeNet第3层：池化层（2*2average池化接relu）
        :param input: 第2层的输出n×28*28*6
        :return: 返回n×14*14*6的张量
        """
        output = self.create_pool(input)
        # print(output)  # shape=(?, 14, 14, 6)
        return output

    def create_c3(self, input):
        """
        创建LeNet第4层：卷积层
        :param input: 第3层的输出n×14*14*6
        :return: 返回n×10*10*16的张量
        """
        # 保存卷积结果
        convs = []

        # 前6个卷积（5*5卷积核，3通道，1输出）
        for i in range(6):
            conv = self.create_conv(
                [5, 5, 3, 1],
                self.channels_slice(input, i, 3)
            )
            convs.append(conv)

        # 后6个卷积（5*5卷积核，4通道，1输出）
        for i in range(6):
            conv = self.create_conv(
                [5, 5, 4, 1],
                self.channels_slice(input, i, 4)
            )
            convs.append(conv)

        # 3个卷积（5*5卷积核，(2+2)通道，1输出）
        for i in range(3):
            a = self.channels_slice(input, i, 2)
            b = self.channels_slice(input, i + 3, 2)
            c = tf.concat((a, b), 3)
            conv = self.create_conv([5, 5, 4, 1], c)
            convs.append(conv)

        # 最后一个卷积（5*5卷积核，6通道，1输出）
        conv = self.create_conv([5, 5, 6, 1], input)
        convs.append(conv)

        # 计算输出
        output = tf.concat(convs, 3)
        # print(output)  # shape=(?, 10, 10, 16)

        # 返回输出
        return output

    def create_s4(self, input):
        """
        创建LeNet第5层：池化层（2*2average池化接relu）
        :param input: 第4层的输出n×10*10*16
        :return: 返回n×5*5*16的张量
        """
        output = self.create_pool(input)
        # print(output)  # shape=(?, 5, 5, 16)
        return output

    def create_c5(self, input):
        """
        创建LeNet第6层：卷积层（16通道，120个卷积核）
        :param input: 第5层的输出n×5*5*16
        :return: 返回n×1*1*120的张量再reshape为n*120
        """
        conv = self.create_conv([5, 5, 16, 120], input)
        output = tf.reshape(conv, [-1, 120])
        # print(output)  # shape=(?, 120)
        return output

    def create_f6(self, input):
        """
        创建LeNet第7层：全连接层(84个神经元)
        :param input: 第6层的输出n×120
        :return: 返回n×84的张量
        """
        output = self.create_fc(input, 120, 84)
        # print(output)  # shape=(?, 84)
        return output

    def create_output(self, input):
        """
        创建LeNet第8层：全连接层(10个神经元)
        :param input: 第7层的输出n×84
        :return: 返回n×10的张量
        """
        fc = self.create_fc(input, 84, 10)
        min = tf.reduce_min(fc)
        max = tf.reduce_max(fc)
        delta = max - min
        output = tf.nn.softmax((fc - min) / delta)
        # print(output)  # shape=(?, 10)
        return output


lenet_cnn = LeNetCNN(
    data_dir='MNIST_data/handwritten_num/',
    model_dir='AI_model/lenet_cnn/'
)

if __name__ == '__main__':
    cnn = lenet_cnn

    # print("begin training...")
    # cnn.train()
    # cnn.save_vars()

    print('\nbegin testing...')
    cnn.eval(cnn.data.test.images, cnn.data.test.labels)

    # while True:
    #     img_path = input('请输入图片路径：')
    #     if not img_path:
    #         break
    #     img = cv2.imread(img_path, 0)
    #     cv2.imshow(img_path, img)
    #     result = cnn.recognize(img)
    #     print('该图片的结果：', result)
    #     cv2.waitKey()
    #     cv2.destroyAllWindows()

    cnn.close()
